#include "LvDescriptors.h"

#include <cassert>
#include <stdexcept>


// *************** Descriptor Set Layout Builder *********************

LvDescriptorSetLayout::Builder& LvDescriptorSetLayout::Builder::addBinding(uint32_t binding,
                                                                             VkDescriptorType descriptorType,
                                                                             VkShaderStageFlags stageFlags,
                                                                             uint32_t count)
{
    assert(bindings.count(binding) == 0 && "Binding already in use");
    VkDescriptorSetLayoutBinding layoutBinding {};
    layoutBinding.binding         = binding;
    layoutBinding.descriptorType  = descriptorType;
    layoutBinding.descriptorCount = count;
    layoutBinding.stageFlags      = stageFlags;
    bindings[binding]             = layoutBinding;
    return *this;
}

std::unique_ptr<LvDescriptorSetLayout> LvDescriptorSetLayout::Builder::build() const
{
    return std::make_unique<LvDescriptorSetLayout>(lvDevice, bindings);
}

// *************** Descriptor Set Layout *********************

LvDescriptorSetLayout::LvDescriptorSetLayout(LvDevice& lvDevice, std::unordered_map<uint32_t, VkDescriptorSetLayoutBinding> bindings) 
    : lvDevice {lvDevice}, bindings {bindings}
{
    std::vector<VkDescriptorSetLayoutBinding> setLayoutBindings {};
    for (auto kv : bindings)
    {
        setLayoutBindings.push_back(kv.second);
    }

    VkDescriptorSetLayoutCreateInfo descriptorSetLayoutInfo {};
    descriptorSetLayoutInfo.sType        = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    descriptorSetLayoutInfo.bindingCount = static_cast<uint32_t>(setLayoutBindings.size());
    descriptorSetLayoutInfo.pBindings    = setLayoutBindings.data();

    if (vkCreateDescriptorSetLayout(lvDevice.device(), &descriptorSetLayoutInfo, nullptr, &descriptorSetLayout) !=
        VK_SUCCESS)
    {
        throw std::runtime_error("failed to create descriptor set layout!");
    }
}

LvDescriptorSetLayout::~LvDescriptorSetLayout()
{
    vkDestroyDescriptorSetLayout(lvDevice.device(), descriptorSetLayout, nullptr);
}

// *************** Descriptor Pool Builder *********************

LvDescriptorPool::Builder& LvDescriptorPool::Builder::addPoolSize(VkDescriptorType descriptorType, uint32_t count)
{
    poolSizes.push_back({descriptorType, count});
    return *this;
}

LvDescriptorPool::Builder& LvDescriptorPool::Builder::setPoolFlags(VkDescriptorPoolCreateFlags flags)
{
    poolFlags = flags;
    return *this;
}
LvDescriptorPool::Builder& LvDescriptorPool::Builder::setMaxSets(uint32_t count)
{
    maxSets = count;
    return *this;
}

std::unique_ptr<LvDescriptorPool> LvDescriptorPool::Builder::build() const
{
    return std::make_unique<LvDescriptorPool>(lvDevice, maxSets, poolFlags, poolSizes);
}

// *************** Descriptor Pool *********************

LvDescriptorPool::LvDescriptorPool(LvDevice& lvDevice,
                                     uint32_t maxSets,
                                     VkDescriptorPoolCreateFlags poolFlags,
                                     const std::vector<VkDescriptorPoolSize>& poolSizes) :
    lvDevice {lvDevice}
{
    VkDescriptorPoolCreateInfo descriptorPoolInfo {};
    descriptorPoolInfo.sType         = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
    descriptorPoolInfo.poolSizeCount = static_cast<uint32_t>(poolSizes.size());
    descriptorPoolInfo.pPoolSizes    = poolSizes.data();
    descriptorPoolInfo.maxSets       = maxSets;
    descriptorPoolInfo.flags         = poolFlags;

    if (vkCreateDescriptorPool(lvDevice.device(), &descriptorPoolInfo, nullptr, &descriptorPool) != VK_SUCCESS)
    {
        throw std::runtime_error("failed to create descriptor pool!");
    }
}

LvDescriptorPool::~LvDescriptorPool() { vkDestroyDescriptorPool(lvDevice.device(), descriptorPool, nullptr); }

bool LvDescriptorPool::allocateDescriptor(const VkDescriptorSetLayout descriptorSetLayout,
                                           VkDescriptorSet& descriptor) const
{
    VkDescriptorSetAllocateInfo allocInfo {};
    allocInfo.sType              = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    allocInfo.descriptorPool     = descriptorPool;
    allocInfo.pSetLayouts        = &descriptorSetLayout;
    allocInfo.descriptorSetCount = 1;

    // Might want to create a "DescriptorPoolManager" class that handles this case, and builds
    // a new pool whenever an old pool fills up. But this is beyond our current scope
    if (vkAllocateDescriptorSets(lvDevice.device(), &allocInfo, &descriptor) != VK_SUCCESS)
    {
        return false;
    }
    return true;
}

void LvDescriptorPool::freeDescriptors(std::vector<VkDescriptorSet>& descriptors) const
{
    vkFreeDescriptorSets(
        lvDevice.device(), descriptorPool, static_cast<uint32_t>(descriptors.size()), descriptors.data());
}

void LvDescriptorPool::resetPool() { vkResetDescriptorPool(lvDevice.device(), descriptorPool, 0); }

// *************** Descriptor Writer *********************

LvDescriptorWriter::LvDescriptorWriter(LvDescriptorSetLayout& setLayout, LvDescriptorPool& pool) :
    setLayout {setLayout}, pool {pool}
{}

LvDescriptorWriter& LvDescriptorWriter::writeBuffer(uint32_t binding, VkDescriptorBufferInfo* bufferInfo)
{
    assert(setLayout.bindings.count(binding) == 1 && "Layout does not contain specified binding");

    auto& bindingDescription = setLayout.bindings[binding];

    assert(bindingDescription.descriptorCount == 1 && "Binding single descriptor info, but binding expects multiple");

    VkWriteDescriptorSet write {};
    write.sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    write.descriptorType  = bindingDescription.descriptorType;
    write.dstBinding      = binding;
    write.pBufferInfo     = bufferInfo;
    write.descriptorCount = 1;

    writes.push_back(write);
    return *this;
}

LvDescriptorWriter& LvDescriptorWriter::writeImage(uint32_t binding, VkDescriptorImageInfo* imageInfo)
{
    assert(setLayout.bindings.count(binding) == 1 && "Layout does not contain specified binding");

    auto& bindingDescription = setLayout.bindings[binding];

    assert(bindingDescription.descriptorCount == 1 && "Binding single descriptor info, but binding expects multiple");

    VkWriteDescriptorSet write {};
    write.sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    write.descriptorType  = bindingDescription.descriptorType;
    write.dstBinding      = binding;
    write.pImageInfo      = imageInfo;
    write.descriptorCount = 1;

    writes.push_back(write);
    return *this;
}

bool LvDescriptorWriter::build(VkDescriptorSet& set)
{
    bool success = pool.allocateDescriptor(setLayout.getDescriptorSetLayout(), set);
    if (!success)
    {
        return false;
    }
    overwrite(set);
    return true;
}

void LvDescriptorWriter::overwrite(VkDescriptorSet& set)
{
    for (auto& write : writes)
    {
        write.dstSet = set;
    }
    vkUpdateDescriptorSets(pool.lvDevice.device(), writes.size(), writes.data(), 0, nullptr);
}
